package com.enzoism.ai.chat;

import com.enzoism.ai.chat.model.ChatRequest;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.memory.InMemoryChatMemory;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.http.MediaType;
import org.springframework.http.codec.ServerSentEvent;
import org.springframework.web.bind.annotation.*;
import reactor.core.publisher.Flux;

import java.util.Map;

@Slf4j
@RestController
public class ChatController {

    @Resource(name = "openAiChatClient")
    private ChatClient chatClient;


    /**
     * 直接返回文本
     */
    @GetMapping("/ai/generate")
    public Map<String, Object> generate(@RequestParam(value = "message", defaultValue = "给我讲一个笑话") String message) {
        // 使用chatClient生成对话
        Prompt prompt = new Prompt(message);
        String res = chatClient.prompt(prompt).call().content();
        return Map.of("generation", StringUtils.defaultString(res));
    }

    /**
     * 返回ChatResponse结构的对象【非SSE版本】
     */
    @GetMapping("/ai/generateStreamChatResponse")
    public Flux<ChatResponse> generateStreamChatResponse(@RequestParam(value = "message", defaultValue = "给我讲一个笑话") String message) {
        // 使用chatClient生成对话
        Prompt prompt = new Prompt(new UserMessage(message));
        return chatClient.prompt(prompt).stream().chatResponse();
    }

    @GetMapping("/ai/generateStreamString")
    public Flux<String> generateStreamString(@RequestParam(value = "message", defaultValue = "给我讲个笑话") String message) {
        // 使用chatClient生成对话
        Prompt prompt = new Prompt(message);
        return chatClient.prompt(prompt).stream().content();
    }

    /**
     * 使用SSE(Server-Sent Events)实现聊天
     *
     * @param chatRequest 用户请求
     * @return 聊天响应
     */
    @PostMapping(value = "/ai/sse", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<ServerSentEvent<ChatResponse>> sse(@RequestBody ChatRequest chatRequest) {
        // 使用chatClient生成对话
        Prompt prompt = new Prompt(chatRequest.getUserText());
        Flux<ServerSentEvent<ChatResponse>> stream = chatClient.prompt(prompt).stream().chatResponse().map(chatResponse -> {
            ServerSentEvent<ChatResponse> sentEvent = ServerSentEvent.builder(chatResponse).event("message").build();
            return sentEvent;
        });
        log.info(stream.toString());
        return stream;
    }


    private static final String DEFAULT_PROMPT = "你是一个搞笑的智能笑话大王助手，请根据用户提问回答！";

    /**
     * 使用SSE(Server-Sent Events)实现聊天-添加Memory记忆
     *
     * @param chatRequest 用户请求
     * @return 聊天响应
     */
    @PostMapping(value = "/ai/sseMemory", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<ServerSentEvent<ChatResponse>> sseMemory(@RequestBody ChatRequest chatRequest) {
        // 使用chatClient生成对话
        ChatMemory chatMemory = new InMemoryChatMemory();
        String sessionId = "1111111111"; // 根据业务自己去生成
        int maxMessages = 10;
        MessageChatMemoryAdvisor advisor = new MessageChatMemoryAdvisor(chatMemory, sessionId, maxMessages);
        Flux<ServerSentEvent<ChatResponse>> stream = chatClient.prompt().user(chatRequest.getUserText()).system(DEFAULT_PROMPT).advisors(advisor).stream().chatResponse().map(chatResponse -> {
            ServerSentEvent<ChatResponse> sentEvent = ServerSentEvent.builder(chatResponse).event("message").build();
            return sentEvent;
        });
        log.info(stream.toString());
        return stream;
    }
}
